import logging
import os
import random
import re
import shutil
import sys
import tempfile
from typing import TYPE_CHECKING, List, Optional, Set, Tuple

import tqdm

from chernc.agents.translator import TranslatorBase
from chernc.compiler.rust.error_messages import CargoMessageCompilerMessage, RustcErrorMessages, RustcErrorSpan
from chernc.constants import DEFAULT_LOGGER_DIR, DEFAULT_LOGGER_FILE_NAME
from chernc.core.cargo_manager import CargoManager
from chernc.core.repair import (
    annotate_lines_for_code,
    extern_use_rest,
    extract_rust_code_from_md,
    extract_toml_code_from_md,
    filter_error_messages,
    format_error_messages,
)
from chernc.core.result import DummySnippetResult, FileRepairResult, ProjectRepairResult, ProjectResult, SnippetResult, UseManager
from chernc.parser.rust_parser import rust_code_chop
from chernc.textdiff.word_matcher import matched_text
from chernc.treediff.diff import index_matched_blocks_by_tree

if TYPE_CHECKING:
    from chernc.core.transpiler import ProjectTranspiler

from chernc.logging import build_logger, build_simple_file_logger, release_logger

logger = build_logger(__name__ + ".project_repair", DEFAULT_LOGGER_FILE_NAME, logger_dir=DEFAULT_LOGGER_DIR)


class CodeFixRule(object):
    def __init__(self, repair_engine: "ProjectRepairEngine") -> None:
        self.repair_engine = repair_engine

    def is_match(self, message: RustcErrorMessages) -> bool:
        return False

    def fix_code(self, message: RustcErrorMessages, project_result: ProjectRepairResult, repair_logger: logging.Logger):
        pass


class RuleCannotFindExternCrate(CodeFixRule):
    def __init__(self, repair_engine: "ProjectRepairEngine") -> None:
        super().__init__(repair_engine)

    def is_match(self, message: RustcErrorMessages) -> bool:
        m = re.match(r"can't find crate for `[^`]*?`", message.message)
        return message.level == "error" and m is not None

    def fix_code(self, message: RustcErrorMessages, project_result: ProjectRepairResult, repair_logger: logging.Logger):
        pass


class RuleUnusedImports(CodeFixRule):
    def __init__(self, repair_engine: "ProjectRepairEngine") -> None:
        super().__init__(repair_engine)

    def is_match(self, message: RustcErrorMessages) -> bool:
        # "unused imports: `atoi`, `c_void`, `free`, `strcmp`, and `strlen`"
        # "unused import: `libc::malloc as libc_malloc`"
        imports_m = re.match(r"unused imports: (?:`[^`]*?`)+(?:, `[^`]*?`)*, and `[^`]*?`", message.message)
        import_m = re.match(r"unused import: `[^`]*?`", message.message)
        return message.level == "warning" and (imports_m is not None or import_m is not None)

    def fix_code(self, message: RustcErrorMessages, project_result: ProjectRepairResult, repair_logger: logging.Logger):
        if len(message.spans) == 0:
            return
        filepath = self.repair_engine._to_project_path(message.spans[0].file_name)

        unused_list = []
        imports_m = re.match(r"unused imports: ((?:`[^`]*?`)+(?:, `[^`]*?`)*), and `([^`]*?)`", message.message)
        if imports_m is not None:
            parts = imports_m.group(1).split(",")
            for part in parts:
                part = part.strip()
                part = part.removeprefix("`").removesuffix("`")
                paths, useds, as_name = UseManager.parse_use_statement(part)
                if as_name is not None:
                    unused_list.append(as_name)
                else:
                    for name in useds:
                        unused_list.append(name)
            unused_list.append(imports_m.group(2).strip())

        import_m = re.match(r"unused import: `([^`]*?)`", message.message)
        if import_m is not None:
            paths, useds, as_name = UseManager.parse_use_statement(import_m.group(1).strip())
            if as_name is not None:
                unused_list.append(as_name)
            else:
                for name in useds:
                    unused_list.append(name)

        for unused_name in unused_list:
            project_result.files[filepath].uses.dont_use(unused_name)


class RuleUnresolvedImports(CodeFixRule):
    def __init__(self, repair_engine: "ProjectRepairEngine") -> None:
        super().__init__(repair_engine)

    def is_match(self, message: RustcErrorMessages) -> bool:
        # "unresolved import `itoa`"
        # "unresolved imports `crate::src::trie::a`, `crate::src::trie::v`"
        import_m = re.match(r"unresolved import `[^`]*?`", message.message)
        imports_m = re.match(r"unresolved imports `[^`]*?`(?:, `[^`]*?`)+", message.message)
        return message.level == "error" and (imports_m is not None or import_m is not None)

    def fix_code(self, message: RustcErrorMessages, project_result: ProjectRepairResult, repair_logger: logging.Logger):
        if len(message.spans) == 0:
            return
        filepath = self.repair_engine._to_project_path(message.spans[0].file_name)

        unresolved_list = []
        imports_m = re.match(r"unresolved imports (`[^`]*?`(?:, `[^`]*?`)+)", message.message)
        if imports_m is not None:
            parts = imports_m.group(1).split(",")
            for part in parts:
                part = part.strip()
                part = part.removeprefix("`").removesuffix("`")
                paths, useds, as_name = UseManager.parse_use_statement(part)
                if as_name is not None:
                    unresolved_list.append(as_name)
                else:
                    for name in useds:
                        unresolved_list.append(name)

        import_m = re.match(r"unresolved import `([^`]*?)`", message.message)
        if import_m is not None:
            paths, useds, as_name = UseManager.parse_use_statement(import_m.group(1).strip())
            if as_name is not None:
                unresolved_list.append(as_name)
            else:
                for name in useds:
                    unresolved_list.append(name)

        for unused_name in unresolved_list:
            project_result.files[filepath].uses.dont_use(unused_name)


class ProjectRepairEngine:
    def __init__(self, project: "ProjectTranspiler", translator: TranslatorBase):
        self.project = project
        self.translator = translator

    def _to_project_path(self, path: str) -> str:
        norm_path = os.path.normpath(path).replace("\\", "/")
        # norm_path = os.path.relpath(norm_path, self.project.config.source_path)
        return norm_path

    def _get_related_files(self, all_spans: List[RustcErrorSpan]) -> List[str]:
        related_files: Set[str] = set()
        for span in all_spans:
            while span.expansion:
                span = span.expansion.span
            related_files.add(self._to_project_path(span.file_name))
        return list(related_files)

    def _get_prompt_from_span(self, project_result: ProjectRepairResult, file_name: str, all_spans: List[RustcErrorSpan]) -> str:
        if file_name not in project_result.files:
            raise Exception(f"File {file_name} is not in file results")
        file_result: FileRepairResult = project_result.files[file_name]

        # prepare file
        annotated_lines: Set[int] = set()
        for message_span in all_spans:
            if self._to_project_path(message_span.file_name) != file_name:
                continue
            annotated_lines.add(message_span.line_start)

        to_fix = annotate_lines_for_code(file_result.text, annotated_lines)
        return to_fix

    async def _single_rules_repair_project(
        self, project_result: ProjectRepairResult, repair_dir: str, repair_logger: logging.Logger
    ):
        rules = (
            RuleCannotFindExternCrate(self),
            RuleUnusedImports(self),
            RuleUnresolvedImports(self),
        )

        def select_rule(msg: RustcErrorMessages) -> Optional[CodeFixRule]:
            for rule in rules:
                if rule.is_match(msg):
                    return rule
            return None

        # build and get error messages
        with CargoManager.from_cargo_config(project_result.cargo, repair_dir) as cargo_manager:
            self.project.write_project(repair_dir, project_result, with_comments=False)
            messages = cargo_manager.build()
            error_messages = [message.message for message in messages if isinstance(message, CargoMessageCompilerMessage)]

            targets: Set[Tuple[str, str]] = set()  # name, src_path
            for message in messages:
                if not isinstance(message, CargoMessageCompilerMessage):
                    continue
                targets.add((message.target.name, message.target.src_path))

            # 先修复一些警告
            for target_name, target_src_path in targets:
                self.project.write_project(repair_dir, project_result, with_comments=False)
                messages = cargo_manager.build()
                error_messages = [
                    message.message
                    for message in messages
                    if isinstance(message, CargoMessageCompilerMessage) and message.target.src_path == target_src_path
                ]
                if len(error_messages) == 0:
                    continue
                for msg in tqdm.tqdm(
                    error_messages, desc=f"Rule-based project repair (Target: {target_name})", unit="fix", file=sys.stdout
                ):
                    rule = select_rule(msg)
                    if rule is None:
                        continue
                    else:
                        rule.fix_code(msg, project_result, repair_logger)

        return project_result

    async def _single_llm_repair_project(
        self, project_result: ProjectRepairResult, repair_dir: str, repair_logger: logging.Logger
    ):
        # build and get error messages
        cargo_manager = CargoManager.from_cargo_config(project_result.cargo, repair_dir)
        messages = cargo_manager.build()
        error_messages = [message.message for message in messages if isinstance(message, CargoMessageCompilerMessage)]
        error_messages = filter_error_messages(error_messages, anonymous=False)
        if len(error_messages) == 0:
            return project_result

        repair_logger.debug(f"=======ALL ERROR MESSAGES(Number: {len(error_messages)})=======")
        repair_logger.debug(format_error_messages(error_messages, True))

        # select message
        msg_selected = random.choice(error_messages)
        repair_logger.debug("=======SELECTED ERR=======")
        repair_logger.debug(format_error_messages([msg_selected], True))

        if len(msg_selected.spans) == 0:
            return project_result

        related_files = self._get_related_files(msg_selected.spans)
        to_fix = ""
        for file_name in related_files:
            prompt_code = self._get_prompt_from_span(
                project_result=project_result, file_name=file_name, all_spans=msg_selected.spans
            )
            to_fix += f"```rust\n// FILENAME: {file_name}\n{prompt_code}\n```\n"

        resp = await self.translator.fix(
            to_fix=to_fix,
            deps="",
            err_msgs=format_error_messages([msg_selected]),
            temperature=random.randint(0, 10) / 10,  # 随机一个温度，避免模型每次重试输出过于一致
        )

        repair_logger.debug("=======LLM RAW OUTPUT=======")
        repair_logger.debug(resp)

        # TODO: 可能回复多段代码，在多个文件内
        resp_code = extract_rust_code_from_md(resp)
        if toml := extract_toml_code_from_md(resp):
            await cargo_manager.add_from_config(toml)

        repair_logger.debug("=======CODE EXTRACTION=======")
        repair_logger.debug(resp_code)

        # 没有结果
        if not resp_code:
            logger.debug(f"Retry for project because of no result.")
            return project_result

        # 对LLM输出进行切片
        extern_decls, extern_uses, the_rest = extern_use_rest(resp_code)
        fix_slices = rust_code_chop(the_rest)

        # 对于所有文件进行切片
        project_slices: List[Tuple[str, SnippetResult, str]] = []  # (filepath, snippet, chop_slice)
        for file_path in related_files:
            for snippet in project_result.files[file_path].snippets:
                if isinstance(snippet, DummySnippetResult):
                    continue
                original_slices = rust_code_chop(snippet.text)
                for s in original_slices:
                    project_slices.append((file_path, snippet, s))

        # 现在fix_code不是最终结果，因为LLM的回复是不完整的
        # 如果结果是一些top-level片段，我们对每个片段进行替换
        # 我们只能假设这个结果只有top-level片段
        repair_logger.debug(f"Fix slices number: {len(fix_slices)}")
        if fix_slices:
            for be_replaced_i, replacement_i in index_matched_blocks_by_tree([s[2] for s in project_slices], fix_slices):
                file_path = project_slices[be_replaced_i][0]
                related_snippet = project_slices[be_replaced_i][1]

                be_replaced = project_slices[be_replaced_i][2]
                replacement = fix_slices[replacement_i]

                if be_replaced == replacement:
                    continue
                repair_logger.debug(f"=======REPLACED(TOP-LEVEL)======")
                repair_logger.debug(be_replaced)
                repair_logger.debug(f"=======REPLACEMENT(TOP-LEVEL)======")
                repair_logger.debug(replacement)

                related_snippet.text = related_snippet.text.replace(be_replaced, replacement)
        else:
            # 如果结果不是top-level片段，那么就是某个bug行附近的代码
            be_replaced = None
            be_replaced_snippet = None
            for file_path in related_files:
                file_result = project_result.files[file_path]
                file_spans = [span for span in msg_selected.spans if span.file_name == file_path]
                for snippet in file_result.snippets:
                    original_slices = rust_code_chop(snippet.text)
                    annotated_lines = set()
                    for span in file_spans:
                        found_snippet = file_result.find_snippet_by_line(span.line_start - 1)
                        if found_snippet == snippet:
                            annotated_lines.add(span.line_start)

                    be_replaced = matched_text(snippet.text, the_rest, {l - 1 for l in annotated_lines})
                    be_replaced_snippet = snippet
                    if be_replaced is not None:
                        break
                if be_replaced is not None:
                    break

            if be_replaced is None or be_replaced_snippet is None:
                repair_logger.error("Text replacement failed.")
                return project_result

            repair_logger.debug(f"=======REPLACED(TEXT-MATCH)======")
            repair_logger.debug(be_replaced)
            repair_logger.debug(f"=======REPLACEMENT(TEXT-MATCH)======")
            repair_logger.debug(the_rest)
            be_replaced_snippet.text = be_replaced_snippet.text.replace(be_replaced, the_rest)

    async def _single_repair_project(self, project_result: ProjectRepairResult, repair_dir: str) -> ProjectRepairResult:
        repair_logger = build_simple_file_logger(
            f"chernc.{self.project.project_hash}.repair",
            os.path.join(repair_dir, f"project_repair.log"),
        )
        self.project.write_project(repair_dir, project_result, with_comments=False)

        await self._single_rules_repair_project(project_result=project_result, repair_dir=repair_dir, repair_logger=repair_logger)
        await self._single_llm_repair_project(project_result=project_result, repair_dir=repair_dir, repair_logger=repair_logger)
        await self._single_rules_repair_project(project_result=project_result, repair_dir=repair_dir, repair_logger=repair_logger)

        release_logger(repair_logger)
        return project_result

    async def _repair_project_internal(
        self, project_result: ProjectResult, temp_dir: str, max_fix_count: int = 25, max_fix_earlystop: int = 5
    ) -> ProjectRepairResult:
        # Create a logger for repair
        project_fix_result = project_result.to_project_repair_result()
        # write to dir
        repair_attempt_dir = os.path.join(temp_dir, "attempt_0")
        self.project.write_project(repair_attempt_dir, project_fix_result, with_comments=False)

        error_nums = []
        consecutive_no_improvement = 0
        for i in range(max_fix_count):
            last_repair_attempt_dir = os.path.join(temp_dir, f"attempt_{i}")
            repair_attempt_dir = os.path.join(temp_dir, f"attempt_{i+1}")
            repair_log_path = os.path.join(repair_attempt_dir, "log.txt")

            if os.path.exists(repair_attempt_dir):
                shutil.rmtree(repair_attempt_dir)
            project_fix_result = await self._single_repair_project(project_fix_result, repair_attempt_dir)

            # Check errors
            cargo_manager = CargoManager.from_cargo_config(project_result.cargo, repair_attempt_dir)
            messages = cargo_manager.build()
            error_messages = [message.message for message in messages if isinstance(message, CargoMessageCompilerMessage)]
            error_messages = filter_error_messages(error_messages, anonymous=False)

            current_error_num = len(error_messages)
            error_nums.append(current_error_num)
            if len(error_messages) == 0:
                break

            # Check for consecutive no improvement or increase in error count
            if i > 0 and current_error_num >= error_nums[-2]:
                consecutive_no_improvement += 1
            else:
                consecutive_no_improvement = 0

            if consecutive_no_improvement >= max_fix_earlystop:
                break

        return project_fix_result

    async def repair_project(
        self, project_result: ProjectResult, temp_dir: Optional[str] = None, max_fix_count: int = 25
    ) -> ProjectRepairResult:
        logger.info(f"Repairing project to the path {temp_dir}...")
        if temp_dir is None:
            with tempfile.TemporaryDirectory() as temp_dir:
                project_fix_result = await self._repair_project_internal(project_result, temp_dir, max_fix_count)
        else:
            project_fix_result = await self._repair_project_internal(project_result, temp_dir, max_fix_count)

        return project_fix_result
